%% Load the data - Select the File: 'RotationalDiffusion'
close all; clear all; clc;
mag = 3; Camera_um_pix = 3.45; %FLIR
% Load Data
Exp = '0.0';fps_M1 = 15; MinSpeed_M1 = 92.2 - 38.0;
load('.\RotationalDiffusion\M1\ParticleTracks'); tracks_M1 = tracks; tracks = []; 
Exp = '0.5';fps_5cP = 10; MinSpeed_5cP = 21.8 - 6.8;
load('.\RotationalDiffusion\5cP\ParticleTracks'); tracks_5cP = tracks;
%% Filter Tracks Using Length and Min. Swimming speed: 
% M1 Media
clc;
Tracks_sel_M1 = []; tracks_c_M1 = []; TrackLength_M1 = []; c = 0; MinTrackLength_M1 = fps_M1*15; % > 15 seconds
[tracks_c_M1] = Camera2MicronSecond(tracks_M1 ,mag,fps_M1,Camera_um_pix);
for ii = 1:length(tracks_c_M1)
   if length(tracks_c_M1(ii).x) >= MinTrackLength_M1 && mean(tracks_c_M1(ii).Speed_um) > MinSpeed_M1
       c = c + 1;
       Tracks_sel_M1(c).t = tracks_c_M1(ii).t;
       Tracks_sel_M1(c).x = tracks_c_M1(ii).x; Tracks_sel_M1(c).x_um = tracks_c_M1(ii).x_um;
       Tracks_sel_M1(c).y = tracks_c_M1(ii).y; Tracks_sel_M1(c).y_um = tracks_c_M1(ii).y_um;
       Tracks_sel_M1(c).u = tracks_c_M1(ii).u; Tracks_sel_M1(c).u_um = tracks_c_M1(ii).u_um;
       Tracks_sel_M1(c).v = tracks_c_M1(ii).v; Tracks_sel_M1(c).v_um = tracks_c_M1(ii).v_um;
       Tracks_sel_M1(c).magXY = sqrt((tracks_c_M1(ii).x_um).^2+(tracks_c_M1(ii).y_um).^2);
       Tracks_sel_M1(c).Speed_um = tracks_c_M1(ii).Speed_um; Tracks_sel_M1(c).MeanSpeed_um = mean(tracks_c_M1(ii).Speed_um);
   end
end
myparameters.lengthdev = 4;
[tracks_filt_M1, trackstats] = FilterTracks(Tracks_sel_M1, myparameters);
for i = 1:length(tracks_filt_M1)
    TrackLength_M1(i,1) = length(tracks_filt_M1(i).x);
end
% 5 cP Media
Tracks_sel_5cP = []; tracks_c_5cP = []; TrackLength_5cP = []; c = 0; MinTrackLength_5cP = fps_5cP*15; % > 15 seconds
[tracks_c_5cP] = Camera2MicronSecond(tracks_5cP ,mag,fps_5cP,Camera_um_pix);
for ii = 1:length(tracks_c_5cP)
   if length(tracks_c_5cP(ii).x) >= MinTrackLength_5cP && mean(tracks_c_5cP(ii).Speed_um) > MinSpeed_5cP
       c = c + 1;
       Tracks_sel_5cP(c).t = tracks_c_5cP(ii).t;
       Tracks_sel_5cP(c).x = tracks_c_5cP(ii).x; Tracks_sel_5cP(c).x_um = tracks_c_5cP(ii).x_um;
       Tracks_sel_5cP(c).y = tracks_c_5cP(ii).y; Tracks_sel_5cP(c).y_um = tracks_c_5cP(ii).y_um;
       Tracks_sel_5cP(c).u = tracks_c_5cP(ii).u; Tracks_sel_5cP(c).u_um = tracks_c_5cP(ii).u_um;
       Tracks_sel_5cP(c).v = tracks_c_5cP(ii).v; Tracks_sel_5cP(c).v_um = tracks_c_5cP(ii).v_um;
       Tracks_sel_5cP(c).magXY = sqrt((tracks_c_5cP(ii).x_um).^2+(tracks_c_5cP(ii).y_um).^2);
       Tracks_sel_5cP(c).Speed_um = tracks_c_5cP(ii).Speed_um;
       Tracks_sel_5cP(c).MeanSpeed_um = mean(tracks_c_5cP(ii).Speed_um);
   end
end
myparameters.lengthdev = 4;
[tracks_filt_5cP, trackstats] = FilterTracks(Tracks_sel_5cP, myparameters);
for i = 1:length(tracks_filt_5cP)
    TrackLength_5cP(i,1) = length(tracks_filt_5cP(i).x);
end
'Done Filter'
% Calculate MSD
timeEnd = 20;
[MSD_x_M1, MSD_x_adj_M1, MSD_y_M1, MSD_y_adj_M1, MSD_XY_M1,  MSD_XY_adj_M1, NumberDataPoints_M1] = D_rot_MSD(tracks_filt_M1, fps_M1, timeEnd);
[MSD_x_5cP, MSD_x_adj_5cP, MSD_y_5cP, MSD_y_adj_5cP, MSD_XY_5cP,  MSD_XY_adj_5cP, NumberDataPoints_5cP] = D_rot_MSD(tracks_filt_5cP, fps_5cP, timeEnd);
'Done MSD'
% Calculate Rotational Diffusion
min_val = 1; t_range = linspace(10,20,11);
timeCalc = 15; % Time point to return D_r
% M1
t_M1 = 1:max(TrackLength_M1); t_s_M1 = t_M1./fps_M1;
V_M1 = mean([tracks_filt_M1(:).MeanSpeed_um]); % Mean Swimming Speed
[D_r_M1, tau_o_M1, D_r_15_M1, tau_0_15_M1, modelfun_M1] = D_rot(MSD_XY_M1, t_range, min_val, timeCalc, fps_M1, t_s_M1, V_M1);
xx_15_M1 = t_s_M1(min_val:fps_M1*20); 
y_fit_M1 = modelfun_M1(D_r_15_M1,xx_15_M1); x_fit_M1 = xx_15_M1;

% 5 cP
t_5cP = 1:max(TrackLength_M1); t_s_5cP = t_5cP./fps_5cP;
V_5cP = mean([tracks_filt_5cP(:).MeanSpeed_um]); % Mean Swimming Speed
[D_r_5cP, tau_o_5cP, D_r_15_5cP, tau_0_15_5cP, modelfun_5cP] = D_rot(MSD_XY_5cP, t_range, min_val, timeCalc, fps_5cP, t_s_5cP, V_5cP);
xx_15_5cP = t_s_5cP(min_val:fps_5cP*20); 
y_fit_5cP = modelfun_5cP(D_r_15_5cP,xx_15_5cP); x_fit_5cP = xx_15_5cP;

D_rots_comb = cat(1,D_r_M1,D_r_5cP)
Exp_STD = std(D_rots_comb)
Exp_Mean = round(mean([D_r_15_M1,D_r_15_5cP]),3)
Final_15 = [round(D_r_15_M1,3) round(D_r_15_5cP,3)]
D_r_Sim = mean(Final_15)
% Sensitivity Analysis: varying fit window
Control(1,1) = round(mean(D_r_M1),3);
Control(1,2) = round(std(D_r_M1),3)
Five(1,1) = round(mean(D_r_5cP),3);
Five(1,2) = round(std(D_r_5cP),3)
% Plot
mypapersize = [3.42 3];
PerPEOColor = [0 0 0;
                0.635 0.078 0.184; 
                0.929 0.694 0.125; 
                0.466 0.674 0.188;
                0.85 0.325 0.0980];
LineWidth = 1; markersize = 5; afs = 7;
alw = 1.0; % Border line width
Control_range = 1:8:fps_M1*15.5; Five_range = 1:5:fps_5cP*15.1;

D_rot_CalcPlot = figure('color',[1 1 1]);
    f = plot(x_fit_M1(1:fps_M1*20),y_fit_M1,'-','color',PerPEOColor(1,:),'color',PerPEOColor(1,:),...
            'LineWidth',1);
    hold on;
    p(1) = plot(t_s_M1(Control_range),MSD_XY_M1(Control_range),'LineStyle', 'none',...
        'marker','o','MarkerSize',4,'color',PerPEOColor(1,:),...
        'LineWidth',LineWidth,'MarkerFaceColor',[1 1 1]);
    
    f = plot(x_fit_5cP,y_fit_5cP,'-','color',PerPEOColor(4,:),'color',PerPEOColor(4,:),...
            'LineWidth',1);
    p(2) = plot(t_s_5cP(Five_range),MSD_XY_5cP(Five_range),'color',PerPEOColor(4,:),'LineStyle', 'none',...
        'marker','o','MarkerSize',4,'color',PerPEOColor(4,:),...
        'LineWidth',LineWidth,'MarkerFaceColor',[1 1 1]);
    set(get(get(f(1:end),'Annotation'),'LegendInformation'),'IconDisplayStyle','off')
    lgd = legend(p(:),'0.996 cP','4.48 cP','location','northwest');
    hold off;box on;
   
    ax = gca; ax.FontSize = afs;
    xlim([0 15.5]);ylim([0 7.5E5]); xticks([0,5,10,15]); yticks([0,2E5,4E5,6E5]);
    xlabel('Time t (s)','FontSize',8);ylabel('MSD \langle\Deltar^2\rangle (\mum^2)','FontSize',8);
    
    lgd.EdgeColor = [1 1 1]; lgd.FontName = 'Arial'; lgd.FontSize = afs;
    set(gca, 'FontName', 'Arial'); set(gca,'LineWidth',alw); xlabel(''); ylabel('');  
    set(gca,'XColor','k');set(gca,'YColor','k');set(gca,'LineWidth',alw)
    set(gcf, 'PaperUnits', 'inches', 'papersize', mypapersize, 'PaperPosition', [0 0 mypapersize]);
    set(D_rot_CalcPlot,'Units','inches','Position',[0 0 mypapersize]);

% Save Figure
figdir = '.\ExtendedDataFig7\';
print(D_rot_CalcPlot, '-painters','-dpdf', '-r1200', [figdir,'D_r_Control_5cP.pdf']); 
% Save Data
tab = 1;
dir_Excel = '.\Source data\ExtendedData\';
% M1 Experiment
writematrix(t_s_M1(Control_range)',[dir_Excel,'ExtDataFigure7.xlsx'],'Sheet',tab,'Range','A4')
writematrix(MSD_XY_M1(Control_range)',[dir_Excel,'ExtDataFigure7.xlsx'],'Sheet',tab,'Range','B4')
% M1 Fit
writematrix(x_fit_M1(1:fps_M1*20)',[dir_Excel,'ExtDataFigure7.xlsx'],'Sheet',tab,'Range','C4')
writematrix(y_fit_M1',[dir_Excel,'ExtDataFigure7.xlsx'],'Sheet',tab,'Range','D4')
% 5 cP Experiment
writematrix(t_s_5cP(Five_range)',[dir_Excel,'ExtDataFigure7.xlsx'],'Sheet',tab,'Range','F4')
writematrix(MSD_XY_5cP(Five_range)',[dir_Excel,'ExtDataFigure7.xlsx'],'Sheet',tab,'Range','G4')
% 5 cP Fit
writematrix(x_fit_5cP',[dir_Excel,'ExtDataFigure7.xlsx'],'Sheet',tab,'Range','H4')
writematrix(y_fit_5cP',[dir_Excel,'ExtDataFigure7.xlsx'],'Sheet',tab,'Range','I4')
%% FUNCTION: Convert pix/frame to um/s
function [c_tracks] = Camera2MicronSecond(tracks,mag,fps,Camera_micron_pix)
% Camera2MicronSecond: Converts camera units (frames and pix) to physical
% units (microns and seconds). Cameras have different pixel dimensions. 
% For example, the Zyla has 6.5 microns per pixel
    c_tracks = tracks;
    for i=1:length(c_tracks)
        c_tracks(i).t_sec = c_tracks(i).t/fps;
        c_tracks(i).x_um = c_tracks(i).x.*Camera_micron_pix/mag; c_tracks(i).y_um = c_tracks(i).y.*Camera_micron_pix/mag;
        c_tracks(i).u_um = c_tracks(i).u.*Camera_micron_pix.*fps./mag; c_tracks(i).v_um = c_tracks(i).v.*Camera_micron_pix.*fps./mag;    
        % Calculate the Speed of the Trajectory (magnitude of u-velocity and v-velocity)    
        s = sqrt(c_tracks(i).v.^2+c_tracks(i).u.^2); c_tracks(i).Speed = s;
        s_um = sqrt(c_tracks(i).v_um.^2+c_tracks(i).u_um.^2);  c_tracks(i).Speed_um = s_um;   
    end
end
%% FUNCTION: Calculate MSD
function [MSD_x, MSD_x_adj, MSD_y, MSD_y_adj, MSD_XY,  MSD_XY_adj, NumberDataPoints] = D_rot_MSD(tracks, fps, timeEnd)
    g = 1; x_sub_store = []; y_sub_store = []; magXY_sub_store = []; Delta_X = []; Delta_Y = []; Delta_magXY = [];
    maxTime_MSD = fps*timeEnd;
    for t = 1:maxTime_MSD
        c = 0;
        for k = 1:length(tracks)
            c = c + 1;
            x_sub_store = tracks(k).x_um -circshift(tracks(k).x_um,t); Delta_X{c} = x_sub_store(t+1:t:end);        
            y_sub_store = tracks(k).y_um -circshift(tracks(k).y_um,t); Delta_Y{c} = y_sub_store(t+1:t:end);
            magXY_sub_store = tracks(k).magXY -circshift(tracks(k).magXY,t); Delta_magXY{c} = magXY_sub_store(t+1:t:end);
        end
        x2_sub_bin(g,t) = mean(vertcat(Delta_X{:}).^2);x2_sub_bin_var(g,t) = var(vertcat(Delta_X{:}));
        y2_sub_bin(g,t) = mean(vertcat(Delta_Y{:}).^2);y2_sub_bin_var(g,t) =  var(vertcat(Delta_Y{:}));   
        magXY2_sub_bin(g,t) = mean(vertcat(Delta_X{:}).^2 + vertcat(Delta_Y{:}).^2); magXY2_bin_var(g,t) =  var(vertcat(Delta_X{:}).^2 + vertcat(Delta_Y{:}).^2);
        NumberDataPoint(g,t) = length(vertcat(Delta_X{:}));
        clear x_sub_store y_sub_store magXY_store Delta_X Delta_Y Delta_magXY
    end
    MSD_x = x2_sub_bin; MSD_x_adj = x2_sub_bin_var; MSD_y = y2_sub_bin; MSD_y_adj = y2_sub_bin_var;
    MSD_XY = x2_sub_bin + y2_sub_bin; MSD_XY_adj = y2_sub_bin_var + y2_sub_bin_var;
    NumberDataPoints = NumberDataPoint;
end
%% FUNCTION: Calculate Rotational Diffusion
function [D_r, tau_o, D_r_15, tau_0_15, modelfun] = D_rot(MSD_XY, t_range, min_val, timeCalc, fps, t_s, V)
    D_r = [];tau_o = [];
    for ii = 1:length(t_range)
        maxInd_fit = []; xx = []; yy = []; beta = [];
        maxInd_fit = fps*t_range(ii); 
        xx = t_s(min_val:maxInd_fit); yy = MSD_XY(min_val:maxInd_fit);
        modelfun = @(D_r,xx)(0.5.*V^2./D_r.^2.*(2.*D_r.*xx + exp(-2*D_r.*xx) - 1)); 
        beta0 = [0.08]; beta = nlinfit(xx,yy,modelfun,beta0);
        D_r(ii,1) = beta(1); % Rotational Diffusion
        tau_o(ii,1) = 1/(2*beta(1));
        if t_range(ii) == timeCalc
            D_r_15 = beta(1); tau_0_15 = 1/(2*beta(1));
        end
    end
end